package com.yc.cloud.openai.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.yc.cloud.openai.dto.request.aigc.ImageGenerateParam;
import com.yc.cloud.openai.dto.request.aigc.ImageMJProcessParam;
import com.yc.cloud.openai.dto.resp.aigc.StableDiffusionGenerateResp;
import com.yc.cloud.openai.dto.resp.aigc.StableDiffusionInfoDto;
import com.yc.cloud.openai.service.ImageAiGCService;
import com.yc.cloud.openai.vo.ImageVo;
import lombok.extern.slf4j.Slf4j;
import lombok.val;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

@Service("StableDiffusionServiceImpl")
@Slf4j
public class StableDiffusionServiceImpl implements ImageAiGCService {

    String url = "http://127.0.0.1:7860/sdapi/v1/txt2img"; // API endpoint

    @Override
    public String generate(String prompt, String webhookUrl) {
        JSONObject data = new JSONObject();
        data.putOpt("prompt", prompt);
        if (webhookUrl != null && !webhookUrl.isEmpty()) {
            data.putOpt("callback_url", webhookUrl);
        }
        HttpRequest request = HttpUtil.createPost(url)
                .body(JSONUtil.toJsonStr(data))
                .contentType("application/json");
        try (HttpResponse response = request.execute()) {
            val result = response.body();
            log.info("生成图片结果：{}", result);
            return result;
        } catch (Exception e) {
            e.printStackTrace();
            log.error(e.getMessage());
        }
        return null;
    }

    @Override
    public ImageVo process(ImageMJProcessParam param) {
        return null;
    }

    @Override
    public List<ImageVo> generate(ImageGenerateParam param) {
        JSONObject data = new JSONObject();
        data.putOpt("prompt", param.getPrompt());
        data.putOpt("negative_prompt", "low quality, blurry, distorted");
        data.putOpt("steps", 30);
        data.putOpt("cfg_scale", 7.5);
        data.putOpt("width", param.getWidth());
        data.putOpt("height", param.getHeight());
        data.putOpt("sampler_name", "Euler");
        data.putOpt("enable_hr", false);
        data.putOpt("hr_scale", 2);
//        if (webhookUrl != null && !webhookUrl.isEmpty()) {
//            data.putOpt("callback_url", webhookUrl);
//        }
        HttpRequest request = HttpUtil.createPost(url)
                .body(JSONUtil.toJsonStr(data))
                .contentType("application/json");
        val imageVos = new ArrayList<ImageVo>();
        try (HttpResponse response = request.execute()) {
            val result = response.body();
            log.info("生成图片结果：{}", result);
            if (StrUtil.isNotEmpty(result) && JSONUtil.isJson(result)) {
                val responseJson = JSONUtil.toBean(result, StableDiffusionGenerateResp.class);
                val images = responseJson.getImages();
                if (CollUtil.isNotEmpty(images)) {
                    images.forEach(imageBase64 -> {
                        ImageVo imageVo = new ImageVo();
                        imageVo.setData(imageBase64);
                        val info = responseJson.getInfo();
                        if (StrUtil.isNotEmpty(info)) {
                            val stableDiffusionInfoDto = JSONUtil.toBean(info, StableDiffusionInfoDto.class);
                            imageVo.setWidth(stableDiffusionInfoDto.getWidth());
                            imageVo.setHeight(stableDiffusionInfoDto.getHeight());
                        }

                        imageVos.add(imageVo);
                    });
                    log.info("生成图片数量：{}", images.size());
                }
            }

        } catch (Exception e) {
            e.printStackTrace();
            log.error(e.getMessage());
        }
        return imageVos;
    }


    public static void main(String[] args) {


    }
}
